#include "rtmp_media_source.hpp"
#include "protocol/flv/flv_tag.hpp"
#include "rtmp_media_sink.hpp"

#include "base/thread/thread_worker.hpp"
#include "core/stream_session.hpp"

#include "server/rtmp/rtmp_protocol/rtmp_define.hpp"
#include "server/rtmp/rtmp_message/rtmp_message.hpp"
#include "server/rtmp/rtmp_message/data_message/rtmp_metadata_message.hpp"
#include "transformat/bridge_factory.hpp"
#include "transformat/media_bridge.hpp"
#include "spdlog/spdlog.h"

using namespace mms;

RtmpMediaSource::RtmpMediaSource(ThreadWorker *worker, std::weak_ptr<StreamSession> s) : MediaSource("rtmp", s, worker), av_pkts_(2048), keyframe_indexes_(200)
{
}

bool RtmpMediaSource::init()
{
    return true;
}

RtmpMediaSource::~RtmpMediaSource()
{
    spdlog::info("RtmpMediaSource::~RtmpMediaSource");
}

bool RtmpMediaSource::add_media_sink(std::shared_ptr<MediaSink> media_sink)
{
    MediaSource::add_media_sink(media_sink);
    return true;
}

bool RtmpMediaSource::on_metadata(std::shared_ptr<RtmpMessage> metadata_pkt)
{
    metadata_ = std::make_shared<RtmpMetaDataMessage>();
    int ret = metadata_->decode(metadata_pkt);
    if (ret <= 0)
    {
        spdlog::error("metadata decode failed. ret={}", ret);
        return false;
    }
    has_video_ = metadata_->has_video();
    has_audio_ = metadata_->has_audio();

    if (has_video_)
    {
        auto video_codec_id = metadata_->get_video_codec_id();
        if (video_codec_id == VideoTagHeader::AVC)
        {
        }
        else if (video_codec_id == VideoTagHeader::HEVC)
        {
        }
        else
        {
            spdlog::error("rtmp:unsupport video codec:{:x}", (int32_t)video_codec_id);
            return false;
        }
    }

    if (has_audio_)
    {
        auto audio_codec_id = metadata_->get_audio_codec_id();
        if (audio_codec_id == AudioTagHeader::AAC)
        {
        }
        else
        {
            spdlog::error("rtmp:unsupport audio codec:{:x}", (int32_t)audio_codec_id);
            return false;
        }
    }
    return true;
}

bool RtmpMediaSource::on_audio_packet(std::shared_ptr<RtmpMessage> &audio_pkt)
{
    latest_frame_index_ = av_pkts_.add_pkt(audio_pkt);
    // 1. 头部解析
    AudioTagHeader header;
    int32_t header_consumed = header.decode(audio_pkt->payload_, audio_pkt->payload_size_);
    if (header_consumed < 0)
    {
        return false;
    }
    // 2. sequence header记录
    bool sequence_header = false;
    if (!audio_ready_ && header.sound_format != AudioTagHeader::AAC)
    {
        audio_ready_ = true;
    }
    else if (header.is_seq_header())
    {
        audio_header_ = audio_pkt;
        sequence_header = true;
        audio_ready_ = true;
    }
    // 3. stream_ready判断
    if (!stream_ready_)
    {
        stream_ready_ = (metadata_ != nullptr) && (has_audio_ ? audio_ready_ : true) && (has_video_ ? video_ready_ : true);
    }

    latest_audio_timestamp_ = audio_pkt->timestamp_;
    if (sequence_header)
    {
        return true;
    }
    // 4. sinks唤醒
    if (latest_frame_index_ <= 300 || latest_frame_index_ % 10 == 0)
    {
        std::lock_guard<std::recursive_mutex> lck(sinks_mtx_);
        for (auto sink : lazy_sinks_)
        {
            sink->wakeup();
        }
    }

    return true;
}

bool RtmpMediaSource::on_video_packet(std::shared_ptr<RtmpMessage> &video_pkt)
{
    latest_frame_index_ = av_pkts_.add_pkt(video_pkt);
    video_key_ind++;
    // 解析头部
    VideoTagHeader header;
    int32_t header_consumed = header.decode(video_pkt->payload_, video_pkt->payload_size_);
    if (header_consumed <= 0)
    {
        return false;
    }

    bool sequence_header = false;
    if (header.is_key_frame() && !header.is_seq_header())
    { // 关键帧索引
        std::unique_lock<std::shared_mutex> wlock(keyframe_indexes_rw_mutex_);
        keyframe_indexes_.push_back(latest_frame_index_);
        //spdlog::info("get a KEY frame");
       
        if(last_key_ind){
            //spdlog::info("The video gop size is {}", video_key_ind - last_key_ind);
        }
        last_key_ind = video_key_ind;
    }
    else if (header.is_seq_header())
    {
        video_ready_ = true;
        video_header_ = video_pkt;
        sequence_header = true;
        spdlog::info("get video header");
    }

    if (sequence_header)
    {
        return true;
    }

    if (!stream_ready_)
    {
        stream_ready_ = (metadata_ != nullptr) && (has_audio_ ? audio_ready_ : true) && (has_video_ ? video_ready_ : true);
    }

    latest_video_timestamp_ = video_pkt->timestamp_;

    if (latest_frame_index_ <= 300 || latest_frame_index_ % 10 == 0)
    {
        std::lock_guard<std::recursive_mutex> lck(sinks_mtx_);
        for (auto sink : lazy_sinks_)
        {
            sink->wakeup();
        }
    }

    return true;
}

std::vector<std::shared_ptr<RtmpMessage>> RtmpMediaSource::get_pkts(int64_t &last_pkt_index, uint32_t max_count)
{
    std::vector<std::shared_ptr<RtmpMessage>> pkts;
    if (last_pkt_index == -1)
    { // 第一次获取包
        if (!stream_ready_)
        {
            return {};
        }
        // 先送头部信息
        pkts.push_back(metadata_->msg());
        if (has_video_)
        {
            if (video_header_)
            {
                pkts.push_back(video_header_);
            }
        }

        if (has_audio_)
        {
            if (audio_header_)
            {
                pkts.push_back(audio_header_);
            }
        }

        if (has_video_)
        {
            boost::circular_buffer<uint64_t>::reverse_iterator it;
            int64_t start_idx = -1;
            {
                std::shared_lock<std::shared_mutex> rlock(keyframe_indexes_rw_mutex_);
                it = keyframe_indexes_.rbegin();
                while (it != keyframe_indexes_.rend()) // 找到最后一个I帧
                {
                    auto pkt = av_pkts_.get_pkt(*it);
                    if (pkt)
                    {
                        // if (latest_video_timestamp_ - pkt->timestamp_ >= 2000)
                        // {
                        //     start_idx = *it;
                        //     break;
                        // }
                        // it++;
                        start_idx = *it;
                        break;
                    }
                    else
                    {
                        start_idx = *it;
                        break;
                    }
                }
            }

            if (start_idx < 0)
            {
                pkts.clear();
                return pkts;
            }

            uint32_t pkt_count = 0;
            while (start_idx <= latest_frame_index_ && pkt_count < max_count)
            {
                auto pkt = av_pkts_.get_pkt(start_idx);
                if (pkt)
                {
                    pkts.emplace_back(av_pkts_.get_pkt(start_idx));
                    pkt_count++;
                }
                else
                {
                    break;
                }
                start_idx++;
            }
            last_pkt_index = start_idx;
        }
    }
    else
    {
        int64_t start_idx = last_pkt_index;
        uint32_t pkt_count = 0;
        while (start_idx <= latest_frame_index_ && pkt_count < max_count)
        {
            auto t = av_pkts_.get_pkt(start_idx);
            if (t)
            {
                pkts.emplace_back(av_pkts_.get_pkt(start_idx));
                pkt_count++;
            }
            else
            {
                break;
            }
            start_idx++;
        }
        last_pkt_index = start_idx;
    }

    return pkts;
}

void RtmpMediaSource::close()
{
    std::set<std::shared_ptr<MediaSink>> tmp_sinks_;
    {
        std::lock_guard<std::recursive_mutex> lck(sinks_mtx_);
        tmp_sinks_.swap(lazy_sinks_);
    }

    for (auto sink : tmp_sinks_)
    {
        sink->get_worker()->dispatch([sink]()
                                     { sink->close(); });
    }
    tmp_sinks_.clear();
}

std::shared_ptr<MediaBridge> mms::RtmpMediaSource::get_or_create_bridge(const std::string& id, StreamSession* session)
{
   std::unique_lock<std::shared_mutex> lck(bridges_mtx_);
   std::shared_ptr<MediaBridge> bridge;
    auto it = bridges_.find(id);
    if(it != bridges_.end()){
        bridge = it->second;
    }
    if(bridge){
        return bridge;
    }
    bridge = BridgeFactory::create_bridge(worker_,id,session->get_domain(),
    session->get_app_name(),session->get_stream_name());
    bridge->init(std::weak_ptr<MediaSource>(shared_from_this()));

    auto media_sink = bridge->get_media_sink();
    media_sink->set_source(shared_from_this());
    MediaSource::add_media_sink(media_sink);

    auto media_source = bridge->get_media_source();
    media_source->set_source_info(session->get_domain(), session->get_app_name(), session->get_stream_name());

    bridges_.insert(std::pair(id, bridge));
    media_sink->wakeup();
    return bridge;

}
